import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_experimental.synthetic_data import create_data_generation_chain
from langchain_experimental.tabular_synthetic_data.openai import create_openai_data_generator
from langchain_experimental.tabular_synthetic_data.prompts import SYNTHETIC_FEW_SHOT_PREFIX, SYNTHETIC_FEW_SHOT_SUFFIX
from pydantic.main import BaseModel

load_dotenv()


# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0.8
)

chain = create_data_generation_chain(model)

# 生成数据
result = chain({  # 根据提供的关键词，随机生成一句话
    # 'fields': ['蓝色', '黄色', '美人'],
    # 'preferences': {}
    'fields': {'颜色': ['蓝色', '黄色']},
    'preferences': {'style': '让他像一首诗歌一样'}
})

print(result)
# {'fields': {'颜色': ['蓝色', '黄色']}, 'preferences': {'style': '让他像一首诗歌一样'}, 'text': '在晨曦微露的海边，那间小屋被漆成了宁静的蓝色，屋顶则点缀着温暖的黄色，仿佛是大海与阳光共同谱写的一首温柔诗篇，静静地诉说着时光的故事。'}


# 生成一些结构化数据
# https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.prompt.PromptTemplate.html
# https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.few_shot.FewShotPromptTemplate.html

# 1.定义数据模型
class MedicalBilling(BaseModel):
    patient_id: int  # 患者ID，整数类型
    patient_name: str  # 患姓名，字符串类型
    diagnosis_code: str  # 诊断代码，字符串类型
    procedure_code: str  # 程序代码，字符串类型
    total_charge: float  # 总费用，浮点数类型
    insurance_claim_amount: float  # 保险索赔金额，浮点数类型


# 2.提供一些样例数据
examples = [
    {'example': 'Patient ID:123456，Patient Name: 张娜, Diagnosis Code: J20.9, Procedure Code: 99203,'
                'Total charge: $500, Insurance Claim Amount: %450'},
    {'example': 'Patient ID: 789012,Patient Name: 王兴鹏,Diagnosis Code: M54.5, Procedure Code: 99213, '
                'Total charge:$150, Insurance Claim Amount: %135'},
    {'example': 'Patient ID: 345678,Patient Name: 刘晓辉,Diagnosis Code: E11.9, Procedure Code: 99214,'
                'Total Charge: $300, Insurance Claim Amount: %270'},
]

# 3.创建提示模版，指导AI生成符合样例数据的数据

openai_template = PromptTemplate(input_variables=['example'], template='{example}')

prompt_template = FewShotPromptTemplate(
    prefix=SYNTHETIC_FEW_SHOT_PREFIX,
    suffix=SYNTHETIC_FEW_SHOT_SUFFIX,
    examples=examples,
    example_prompt=openai_template,
    input_variables=['subject', 'extra']
)

# 4.创建一个结构化数据的生成器
generator = create_openai_data_generator(
    output_schema=MedicalBilling,
    llm=model,
    prompt=prompt_template
)

# 5.调用生成器
result = generator.generate(
    subject='医疗账单',
    extra='名字可以是随机的，最好是一些比较生僻的名字；医疗总费用呈现正态分布，保险索赔金额为医疗总费用的90%。',
    runs=10  # 生成10条数据
)

print(result)
# [
# MedicalBilling(patient_id=502314, patient_name='许峻峰', diagnosis_code='J45.9', procedure_code='99215', total_charge=450.0, insurance_claim_amount=405.0),
# MedicalBilling(patient_id=617283, patient_name='蔡炎', diagnosis_code='M54.5', procedure_code='99213', total_charge=160.0, insurance_claim_amount=144.0),
# MedicalBilling(patient_id=730291, patient_name='周旭东', diagnosis_code='I20.9', procedure_code='99284', total_charge=350.0, insurance_claim_amount=315.0),
# MedicalBilling(patient_id=846172, patient_name='欧阳辰', diagnosis_code='J45.9', procedure_code='99214', total_charge=300.0, insurance_claim_amount=270.0),
# MedicalBilling(patient_id=957381, patient_name='慕容高飞', diagnosis_code='E11.9', procedure_code='99215', total_charge=400.0, insurance_claim_amount=360.0),
# MedicalBilling(patient_id=648273, patient_name='南宫问天', diagnosis_code='M54.3', procedure_code='99238', total_charge=450.0, insurance_claim_amount=405.0),
# MedicalBilling(patient_id=759382, patient_name='上官婉儿', diagnosis_code='I20.0', procedure_code='99284', total_charge=500.0, insurance_claim_amount=450.0),
# MedicalBilling(patient_id=537894, patient_name='令狐冲', diagnosis_code='J45.9', procedure_code='99214', total_charge=420.0, insurance_claim_amount=378.0),
# MedicalBilling(patient_id=864197, patient_name='慕容云海', diagnosis_code='M54.3', procedure_code='99238', total_charge=460.0, insurance_claim_amount=414.0),
# MedicalBilling(patient_id=347621, patient_name='欧阳星辰', diagnosis_code='I20.0', procedure_code='99284', total_charge=510.0, insurance_claim_amount=459.0)
# ]
